import numpy as np
import scipy.signal as signal

from py3gpp import *
from NR.nr import NR
from NR.ssb import SSBInfo

class NRSync(NR):
    def __init__(self,fn,mu=1,centered=True) -> None:
        NR.__init__(self,fn,mu,centered)
        
    def sync(self):
        # peak_value = np.zeros(3)
        # peak_index = np.zeros(3, 'int')
        # PSS_LEN = 128
        # NRE_PER_PRB = 12
        # wv_cut = int(self.fs*25e-3)
        # wv = self.waveform[:wv_cut]
        # start = 56
        # pssIndices = np.arange(start, start + PSS_LEN  - 1)
        # SampleRate = self.fs

        # for current_NID2 in np.arange(3, dtype='int'):
            
        #     self.slotGrid[pssIndices] = nrPSS(current_NID2)
        #     [refWaveform, info] = nrOFDMModulate(self.carrier, self.slotGrid, SampleRate = SampleRate)

        #     refWaveform = refWaveform[info['CyclicPrefixLengths'][0]:]; # remove CP

        #     temp = signal.correlate(wv, refWaveform, 'valid')  # correlate over 25 ms
        #     index = np.argmax(np.abs(temp))
        #     peak = np.abs(temp[index])
        #     peak_index[current_NID2] = index
        #     peak_value[current_NID2] = peak

        # self.detected_NID2 = np.argmax(peak_value)
        # self.timingOffset = peak_index[self.detected_NID2]
        # self.peak =  peak_value[self.detected_NID2]
        NR.sync(self,SampleRate=self.fs)
        return self.step0()
        
    def find_psses(self,waveform,alpha):
        PSS_LEN = 128
        NRE_PER_PRB = 12
        start = 56
        pssIndices = np.arange(start, start + PSS_LEN  - 1)
        SampleRate = self.fs
        all_pss = []
        for nid2 in range(3):
            self.slotGrid[pssIndices] = nrPSS(nid2)
            [refWaveform, info] = nrOFDMModulate(self.carrier, self.slotGrid, SampleRate = SampleRate)
            refWaveform = refWaveform[info['CyclicPrefixLengths'][0]:]; # remove CP
            temp = signal.correlate(waveform, refWaveform, 'valid')  # correlate over 25 ms       
            index = 10
            while index < len(temp):
                if np.abs(temp[index])>self.peak*alpha:
                    l_ = np.argmax(np.abs(temp[index-5:index+5]))
                    all_pss.append((nid2, index-5+l_,temp[index-5+l_]))
                    index += 10
                else:
                    index += 1
        return all_pss
    
    def find_all_pss(self,alpha):
        self.decimate(self.gscn)
        self.all_pss = self.find_psses(alpha)
        return self.all_pss
    
    def sssCheck(self,time_off=None):
        if time_off==None:
            time_off = self.timingOffset
        self.decimate(self.gscn)
        tos = time_off - (127-118)*4
        print(tos)
        rxGrid = nrOFDMDemodulate(carrier = self.carrier, 
                                  waveform = self.waveform[tos:], 
                                  initialNSlot = 2, SampleRate=self.fs, 
                                  CyclicPrefixFraction=0.5)
        rxGrid /= np.max((rxGrid.real.max(), rxGrid.imag.max()))
        sssEst = np.zeros(336*3)
        start = 56
        PSS_LEN = 128
        sssIndices = np.arange(start, start + PSS_LEN  - 1)
        sssRx = nrExtractResources(sssIndices, rxGrid[:,2:])
        for ncellid in range(336*3):

            sssRef = nrSSS(ncellid)
            sssEst[ncellid] = np.abs(np.vdot(sssRx, sssRef))

        self.detected_NID = np.argmax(sssEst)
        return self.detected_NID
    
    def detect_NID(self,nid2,offset):
        scsSSB = 15 * 2**(self.mu)
        nSlot = 0
        rxSampleRate = self.fs
        nrbSSB = 20

        pss_pos = 592
        rxGrid = nrOFDMDemodulate(waveform = self.waveform[offset-pss_pos:][:np.min((len(self.waveform), 2048*8))], 
                                  nrb = nrbSSB, scs = scsSSB, initialNSlot = nSlot, 
                                  SampleRate=rxSampleRate, CyclicPrefixFraction=0.5)
        rxGrid = rxGrid[:,1:5]
        rxGrid /= np.max((rxGrid.real.max(), rxGrid.imag.max()))
        sssIndices = nrSSSIndices()
        sssRx = nrExtractResources(sssIndices, rxGrid)
        sssRx /= np.max((sssRx.real.max(), sssRx.imag.max())) # scale sssRx symbols individually
        
        sssEst = np.zeros(336)
        for NID1 in range(335):
            ncellid = (3*NID1) + nid2
            sssRef = nrSSS(ncellid)
            sssEst[NID1] = np.abs(np.vdot(sssRx, sssRef))
        detected_NID1 = np.argmax(sssEst)
        peak = sssEst[detected_NID1]
        detected_NID = detected_NID1*3 + nid2
        return detected_NID,peak
        
    def checkSSB(self,nid2,offset):
        return self.checkSSBwithWaveform(self.waveform,nid2,offset)
    
    def checkSSBwithWaveform(self,waveform,nid2,offset):
        scsSSB = 15 * 2**(self.mu)
        nSlot = 0
        rxSampleRate = self.fs
        nrbSSB = 20

        pss_pos = 592 if self.mu==1 else 588
        rxGrid = nrOFDMDemodulate(waveform = waveform[offset-pss_pos:][:np.min((len(waveform), 2048*8))], 
                                  nrb = nrbSSB, scs = scsSSB, initialNSlot = nSlot, 
                                  SampleRate=rxSampleRate, CyclicPrefixFraction=0.5)
        rxGrid = rxGrid[:,1:5]
        rxGrid /= np.max((rxGrid.real.max(), rxGrid.imag.max()))
        sssIndices = nrSSSIndices()
        sssRx = nrExtractResources(sssIndices, rxGrid)
        sssRx /= np.max((sssRx.real.max(), sssRx.imag.max())) # scale sssRx symbols individually
        
        sssEst = np.zeros(336)
        for NID1 in range(335):
            ncellid = (3*NID1) + nid2
            sssRef = nrSSS(ncellid)
            sssEst[NID1] = np.abs(np.vdot(sssRx, sssRef))
        detected_NID1 = np.argmax(sssEst)
        detected_NID = detected_NID1*3 + nid2
        ssb = SSBInfo().cellid(detected_NID).offset(offset).set_peak(sssEst[detected_NID1])
        # print("sss nid",detected_NID)
        dmrsIndices = nrPBCHDMRSIndices(detected_NID, style='matlab')
        xcorrPBCHDMRS = np.empty(7)
        for ibar_SSB in range(7):
            PBCHDMRS = nrPBCHDMRS(detected_NID, ibar_SSB)
            xcorrPBCHDMRS[ibar_SSB] = np.abs(np.vdot(nrExtractResources(dmrsIndices, rxGrid), PBCHDMRS))
        ibar_SSB = np.argmax(np.abs(xcorrPBCHDMRS))
        ssb = ssb.set_ibar(ibar_SSB)
        # print(f"detect {detected_NID} ibar_ssb {ibar_SSB}")
        
        # refGrid = np.zeros((nrbSSB*12, 4), 'complex')
        # dmrsIndices = nrPBCHDMRSIndices(detected_NID)
        # temp = nrSetResources(dmrsIndices, refGrid, nrPBCHDMRS(detected_NID, ibar_SSB))
        # H, nVar = nrChannelEstimate(rxGrid = rxGrid, refGrid = refGrid)
        # pbchIndices = nrPBCHIndices(detected_NID)
        # pbch_eqed, csi = nrEqualizeMMSE(nrExtractResources(pbchIndices, rxGrid), nrExtractResources(pbchIndices, H), nVar)
        # pbchBits = nrSymbolDemodulate(pbch_eqed, 'QPSK', nVar, 'soft')
        # E = 864
        # v = ibar_SSB
        # scrambling_seq = nrPBCHPRBS(detected_NID, v, E)
        # scrambling_seq_bpsk = (-1)*scrambling_seq*2 + 1
        # pbchBits_descrambled = pbchBits * scrambling_seq_bpsk
        # pbchBits_csi = pbchBits_descrambled * np.repeat(csi, 2)
        # decoded3, crc_result, payload, lsbotfsfn, hrf, msbidxoffset = nrBCHDecode(pbchBits_csi,8,0,detected_NID)
        
        crc_result, lsbotfsfn = self.demodulate_pbch(self.waveform,ssb) 
        if crc_result==0:
            self.ssb = ssb
        return crc_result, lsbotfsfn, ssb
    def demodulate_pbch(self,waveform,ssb):
        scsSSB = 15 * 2**(self.mu)
        nSlot = 0
        rxSampleRate = self.fs
        nrbSSB = 20

        pss_pos = 592
        offset = ssb.off
        detected_NID = ssb.nid
        ibar_SSB = ssb.ibar
        
        rxGrid = nrOFDMDemodulate(waveform = waveform[offset-pss_pos:][:np.min((len(waveform), 2048*8))], 
                                  nrb = nrbSSB, scs = scsSSB, initialNSlot = nSlot, 
                                  SampleRate=rxSampleRate, CyclicPrefixFraction=0.5)
        
        rxGrid = rxGrid[:,1:5]
        refGrid = np.zeros((nrbSSB*12, 4), 'complex')
        dmrsIndices = nrPBCHDMRSIndices(detected_NID)
        temp = nrSetResources(dmrsIndices, refGrid, nrPBCHDMRS(detected_NID, ibar_SSB))
        H, nVar = nrChannelEstimate(rxGrid = rxGrid, refGrid = refGrid)
        pbchIndices = nrPBCHIndices(detected_NID)
        pbch_eqed, csi = nrEqualizeMMSE(nrExtractResources(pbchIndices, rxGrid), nrExtractResources(pbchIndices, H), nVar)
        pbchBits = nrSymbolDemodulate(pbch_eqed, 'QPSK', nVar, 'soft')
        E = 864
        v = ibar_SSB
        scrambling_seq = nrPBCHPRBS(detected_NID, v, E)
        scrambling_seq_bpsk = (-1)*scrambling_seq*2 + 1
        pbchBits_descrambled = pbchBits * scrambling_seq_bpsk
        pbchBits_csi = pbchBits_descrambled * np.repeat(csi, 2)
        decoded3, crc_result, payload, lsbotfsfn, hrf, msbidxoffset = nrBCHDecode(pbchBits_csi,8,0,detected_NID)
         
        return crc_result, lsbotfsfn
        
    def checkSSB_all(self):
        self.all_ssb = []
        print(f"checking SSB, total: {len(self.all_pss)}")
        for nid2,offset,_ in self.all_pss:
            crc_result, detected_NID, lsbotfsfn, _, _ = self.checkSSB(nid2,offset)
            if crc_result==0:
                self.all_ssb.append(SSBInfo().cellid(detected_NID).offset(offset-592))
        return self.all_ssb
    
    def step0(self):
        self.frames = int(len(self.waveform)/(self.fs*20e-3))-1
        crc_result, lsbotfsfn, ssb = self.checkSSB(self.detected_NID2,self.timingOffset)
        self.synced = (crc_result==0)
        if self.synced:
            self.detected_NID = ssb.nid
            self.ibar_SSB = ssb.ibar
        return crc_result, lsbotfsfn, ssb
    
    def report_nid(self):
        if self.synced:
            return f" {self.ssb}"
        else:
            return "no synced"
    
    def search_t(self, wv, ref):
        l = len(wv)
        res = 256
        step = 1/res
        fwv = np.fft.fft(wv)
        fref = np.fft.fft(ref)
        est = np.zeros((6*res,))
        for i in range(6*res):
            est[i] = np.abs(np.vdot(fwv*np.exp(1j*2*np.pi*np.arange(l)*(i-3*res)*step/l),fref))
        i = np.argmax(est)
        return (i-3*res)*step, est
    
    def calibration_t(self, wv, t):
        l = len(wv)
        fwv = np.fft.fft(wv)
        ret = np.fft.ifft(fwv*np.exp(1j*2*np.pi*np.arange(l)*t/l))
        return ret
    
    def est(self):
        SampleRate = self.fs
        ss = int(round(20e-3*SampleRate))
        slotGrid = nrResourceGrid(self.carrier)
        slotGrid = slotGrid[:,:4]
        nrSetResources(nrPSSIndices(),slotGrid,nrPSS(self.detected_NID%3))
        nrSetResources(nrSSSIndices(),slotGrid,nrSSS(self.detected_NID))
        dmrsIndices = nrPBCHDMRSIndices(self.detected_NID)
        nrSetResources(dmrsIndices, slotGrid, nrPBCHDMRS(self.detected_NID, self.ibar_SSB))
        [refWaveform, info] = nrOFDMModulate(self.carrier, slotGrid, SampleRate = SampleRate)
        refWaveform = refWaveform[info['CyclicPrefixLengths'][0]:]; # remove CP
        ref_len = len(refWaveform)
        self.phase = np.zeros((self.frames,),dtype=np.complex128)
        self.dt = np.zeros((self.frames,))
        for i in range(self.frames):
            wv = self.waveform[self.timingOffset+i*ss:self.timingOffset+i*ss+ref_len]
            self.dt[i],_ = self.search_t(wv,refWaveform)
            # c_wv = self.calibration_t(wv,self.dt[i])
            c_wv =wv
            self.phase[i] = np.dot(c_wv,refWaveform.conj())

    def step1(self,alpha=1.0):
        SampleRate = self.fs
        ss = int(round(20e-3*SampleRate))
        self.est()
        ssb_len = 512*8
        self.ssbWaveform = np.zeros((2*ssb_len,),dtype=np.complex64)
        for i in range(min(self.frames,self.max_frames)):
            wv = self.waveform[self.timingOffset+i*ss-ssb_len:self.timingOffset+i*ss+ssb_len]
            # c_wv = self.calibration_t(wv,self.dt[i])
            self.ssbWaveform += wv * self.phase[i].conj()
        self.ssbWaveform /= max(self.ssbWaveform.real.max(), self.ssbWaveform.imag.max()) 
        self.step1_pss = self.find_psses(self.ssbWaveform,alpha) 
        self.step1_pss.sort(key=lambda x:np.abs(x[2]))
        return self.step1_pss
    
    def step2(self):
        ssb_len = 512*8
        self.ssbs = []
        self.step2_sss = []
        for nid2, off, _ in self.step1_pss[:10]:
            crc_result, detected_NID, _, ibar_SSB = self.checkSSB(nid2,self.timingOffset+off-ssb_len)
            detected_NID, peak = self.detect_NID(self.ssbWaveform,nid2,off)
            self.step2_sss.append((detected_NID,np.abs(peak)))
        
    def step3(self):
        ssb_len = 512*8
        self.ssbs = []
        for nid2, off, _ in self.step1_pss[:10]:
            crc_result, detected_NID, _, ibar_SSB = self.checkSSB(nid2,self.timingOffset+off-ssb_len)
            detected_NID, peak = self.detect_NID(self.ssbWaveform,nid2,off)
            if crc_result == 0:
                detected_NID, peak = self.detect_NID(self.ssbWaveform,nid2,off)
                self.ssbs.append(SSBInfo().cellid(detected_NID).offset(self.timingOffset+off-ssb_len).set_peak(peak).set_ibar(ibar_SSB))
        return self.ssbs
    
    # def cfo_correct(self,f):
        